{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training Classification"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/chemprop/chemprop/blob/main/examples/training_classification.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install chemprop from GitHub if running in Google Colab\n",
    "import os\n",
    "\n",
    "if os.getenv(\"COLAB_RELEASE_TAG\"):\n",
    "    try:\n",
    "        import chemprop\n",
    "    except ImportError:\n",
    "        !git clone https://github.com/chemprop/chemprop.git\n",
    "        %cd chemprop\n",
    "        !pip install .\n",
    "        %cd examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from pathlib import Path\n",
    "\n",
    "from lightning import pytorch as pl\n",
    "\n",
    "from chemprop import data, featurizers, models, nn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Change data inputs here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "chemprop_dir = Path.cwd().parent\n",
    "input_path = chemprop_dir / \"tests\" / \"data\" / \"classification\" / \"mol.csv\" # path to your data .csv file\n",
    "num_workers = 0 # number of workers for dataloader. 0 means using main process for data loading\n",
    "smiles_column = 'smiles' # name of the column containing SMILES strings\n",
    "target_columns = ['NR-AhR', 'NR-ER', 'SR-ARE', 'SR-MMP'] # classification of activity (either 0 or 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>smiles</th>\n",
       "      <th>NR-AhR</th>\n",
       "      <th>NR-ER</th>\n",
       "      <th>SR-ARE</th>\n",
       "      <th>SR-MMP</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>CCOc1ccc2nc(S(N)(=O)=O)sc2c1</td>\n",
       "      <td>1.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>CCN1C(=O)NC(c2ccccc2)C1=O</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>CC[C@]1(O)CC[C@H]2[C@@H]3CCC4=CCCC[C@@H]4[C@H]...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>CCCN(CC)C(CC)C(=O)Nc1c(C)cccc1C</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>CC(O)(P(=O)(O)O)P(=O)(O)O</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>495</th>\n",
       "      <td>Cc1ccccc1CO[C@H]1C[C@]2(C(C)C)CC[C@@]1(C)O2</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>496</th>\n",
       "      <td>NNc1ccc(C(=O)O)cc1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>497</th>\n",
       "      <td>CCCCCCOc1ccccc1C(=O)O</td>\n",
       "      <td>0.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>498</th>\n",
       "      <td>O=C(OCc1ccccc1)C(=O)OCc1ccccc1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>499</th>\n",
       "      <td>CCCSc1ccc2[nH]c(NC(=O)OC)nc2c1</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>500 rows × 5 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                smiles  NR-AhR  NR-ER  SR-ARE  \\\n",
       "0                         CCOc1ccc2nc(S(N)(=O)=O)sc2c1     1.0    NaN     1.0   \n",
       "1                            CCN1C(=O)NC(c2ccccc2)C1=O     0.0    0.0     NaN   \n",
       "2    CC[C@]1(O)CC[C@H]2[C@@H]3CCC4=CCCC[C@@H]4[C@H]...     NaN    NaN     0.0   \n",
       "3                      CCCN(CC)C(CC)C(=O)Nc1c(C)cccc1C     0.0    0.0     NaN   \n",
       "4                            CC(O)(P(=O)(O)O)P(=O)(O)O     0.0    0.0     0.0   \n",
       "..                                                 ...     ...    ...     ...   \n",
       "495        Cc1ccccc1CO[C@H]1C[C@]2(C(C)C)CC[C@@]1(C)O2     NaN    0.0     0.0   \n",
       "496                                 NNc1ccc(C(=O)O)cc1     NaN    NaN     0.0   \n",
       "497                              CCCCCCOc1ccccc1C(=O)O     0.0    NaN     0.0   \n",
       "498                     O=C(OCc1ccccc1)C(=O)OCc1ccccc1     0.0    0.0     0.0   \n",
       "499                     CCCSc1ccc2[nH]c(NC(=O)OC)nc2c1     1.0    1.0     0.0   \n",
       "\n",
       "     SR-MMP  \n",
       "0       0.0  \n",
       "1       0.0  \n",
       "2       NaN  \n",
       "3       0.0  \n",
       "4       0.0  \n",
       "..      ...  \n",
       "495     0.0  \n",
       "496     0.0  \n",
       "497     0.0  \n",
       "498     0.0  \n",
       "499     1.0  \n",
       "\n",
       "[500 rows x 5 columns]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_input = pd.read_csv(input_path)\n",
    "df_input"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get SMILES and targets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "smis = df_input.loc[:, smiles_column].values\n",
    "ys = df_input.loc[:, target_columns].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array(['CCOc1ccc2nc(S(N)(=O)=O)sc2c1', 'CCN1C(=O)NC(c2ccccc2)C1=O',\n",
       "        'CC[C@]1(O)CC[C@H]2[C@@H]3CCC4=CCCC[C@@H]4[C@H]3CC[C@@]21C',\n",
       "        'CCCN(CC)C(CC)C(=O)Nc1c(C)cccc1C', 'CC(O)(P(=O)(O)O)P(=O)(O)O'],\n",
       "       dtype=object),\n",
       " array([[ 1., nan,  1.,  0.],\n",
       "        [ 0.,  0., nan,  0.],\n",
       "        [nan, nan,  0., nan],\n",
       "        [ 0.,  0., nan,  0.],\n",
       "        [ 0.,  0.,  0.,  0.]]))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Take a look at the first 5 SMILES strings and target columns\n",
    "smis[:5], ys[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get molecule datapoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[09:05:03] WARNING: not removing hydrogen atom without neighbors\n"
     ]
    }
   ],
   "source": [
    "all_data = [data.MoleculeDatapoint.from_smi(smi, y) for smi, y in zip(smis, ys)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Perform data splitting for training, validation, and testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['SCAFFOLD_BALANCED',\n",
       " 'RANDOM_WITH_REPEATED_SMILES',\n",
       " 'RANDOM',\n",
       " 'KENNARD_STONE',\n",
       " 'KMEANS']"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# available split types\n",
    "list(data.SplitType.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "mols = [d.mol for d in all_data]  # RDkit Mol objects are use for structure based splits\n",
    "train_indices, val_indices, test_indices = data.make_split_indices(mols, \"random\", (0.8, 0.1, 0.1))\n",
    "train_data, val_data, test_data = data.split_data_by_indices(\n",
    "    all_data, train_indices, val_indices, test_indices\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get MoleculeDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "featurizer = featurizers.SimpleMoleculeMolGraphFeaturizer()\n",
    "\n",
    "train_dset = data.MoleculeDataset(train_data[0], featurizer)\n",
    "val_dset = data.MoleculeDataset(val_data[0], featurizer)\n",
    "test_dset = data.MoleculeDataset(test_data[0], featurizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = data.build_dataloader(train_dset, num_workers=num_workers)\n",
    "val_loader = data.build_dataloader(val_dset, num_workers=num_workers, shuffle=False)\n",
    "test_loader = data.build_dataloader(test_dset, num_workers=num_workers, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Change Message-Passing Neural Network (MPNN) inputs here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Message Passing\n",
    "A `Message passing` constructs molecular graphs using message passing to learn node-level hidden representations.\n",
    "\n",
    "Options are `mp = nn.BondMessagePassing()` or `mp = nn.AtomMessagePassing()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "mp = nn.BondMessagePassing()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Aggregation\n",
    "An `Aggregation` is responsible for constructing a graph-level representation from the set of node-level representations after message passing.\n",
    "\n",
    "Available options can be found in ` nn.agg.AggregationRegistry`, including\n",
    "- `agg = nn.MeanAggregation()`\n",
    "- `agg = nn.SumAggregation()`\n",
    "- `agg = nn.NormAggregation()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ClassRegistry {\n",
      "    'mean': <class 'chemprop.nn.agg.MeanAggregation'>,\n",
      "    'sum': <class 'chemprop.nn.agg.SumAggregation'>,\n",
      "    'norm': <class 'chemprop.nn.agg.NormAggregation'>\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "print(nn.agg.AggregationRegistry)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "agg = nn.MeanAggregation()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Feed-Forward Network (FFN)\n",
    "\n",
    "A `FFN` takes the aggregated representations and make target predictions.\n",
    "\n",
    "Available options can be found in `nn.PredictorRegistry`.\n",
    "\n",
    "For regression:\n",
    "- `ffn = nn.RegressionFFN()`\n",
    "- `ffn = nn.MveFFN()`\n",
    "- `ffn = nn.EvidentialFFN()`\n",
    "\n",
    "For classification:\n",
    "- `ffn = nn.BinaryClassificationFFN()`\n",
    "- `ffn = nn.BinaryDirichletFFN()`\n",
    "- `ffn = nn.MulticlassClassificationFFN()`\n",
    "- `ffn = nn.MulticlassDirichletFFN()`\n",
    "\n",
    "For spectral:\n",
    "- `ffn = nn.SpectralFFN()` # will be available in future version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ClassRegistry {\n",
      "    'regression': <class 'chemprop.nn.predictors.RegressionFFN'>,\n",
      "    'regression-mve': <class 'chemprop.nn.predictors.MveFFN'>,\n",
      "    'regression-evidential': <class 'chemprop.nn.predictors.EvidentialFFN'>,\n",
      "    'regression-quantile': <class 'chemprop.nn.predictors.QuantileFFN'>,\n",
      "    'classification': <class 'chemprop.nn.predictors.BinaryClassificationFFN'>,\n",
      "    'classification-dirichlet': <class 'chemprop.nn.predictors.BinaryDirichletFFN'>,\n",
      "    'multiclass': <class 'chemprop.nn.predictors.MulticlassClassificationFFN'>,\n",
      "    'multiclass-dirichlet': <class 'chemprop.nn.predictors.MulticlassDirichletFFN'>,\n",
      "    'spectral': <class 'chemprop.nn.predictors.SpectralFFN'>\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "print(nn.PredictorRegistry)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "ffn = nn.BinaryClassificationFFN(n_tasks = len(target_columns))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Batch Norm\n",
    "A `Batch Norm` normalizes the outputs of the aggregation by re-centering and re-scaling.\n",
    "\n",
    "Whether to use batch norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_norm = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Metrics\n",
    "`Metrics` are the ways to evaluate the performance of model predictions.\n",
    "\n",
    "Available options can be found in `metrics.MetricRegistry`, including"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ClassRegistry {\n",
      "    'mse': <class 'chemprop.nn.metrics.MSE'>,\n",
      "    'mae': <class 'chemprop.nn.metrics.MAE'>,\n",
      "    'rmse': <class 'chemprop.nn.metrics.RMSE'>,\n",
      "    'bounded-mse': <class 'chemprop.nn.metrics.BoundedMSE'>,\n",
      "    'bounded-mae': <class 'chemprop.nn.metrics.BoundedMAE'>,\n",
      "    'bounded-rmse': <class 'chemprop.nn.metrics.BoundedRMSE'>,\n",
      "    'r2': <class 'chemprop.nn.metrics.R2Score'>,\n",
      "    'binary-mcc': <class 'chemprop.nn.metrics.BinaryMCCMetric'>,\n",
      "    'multiclass-mcc': <class 'chemprop.nn.metrics.MulticlassMCCMetric'>,\n",
      "    'roc': <class 'chemprop.nn.metrics.BinaryAUROC'>,\n",
      "    'prc': <class 'chemprop.nn.metrics.BinaryAUPRC'>,\n",
      "    'accuracy': <class 'chemprop.nn.metrics.BinaryAccuracy'>,\n",
      "    'f1': <class 'chemprop.nn.metrics.BinaryF1Score'>\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "print(nn.metrics.MetricRegistry)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# AUROC used by default\n",
    "metric_list = None "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constructs MPNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MPNN(\n",
       "  (message_passing): BondMessagePassing(\n",
       "    (W_i): Linear(in_features=86, out_features=300, bias=False)\n",
       "    (W_h): Linear(in_features=300, out_features=300, bias=False)\n",
       "    (W_o): Linear(in_features=372, out_features=300, bias=True)\n",
       "    (dropout): Dropout(p=0.0, inplace=False)\n",
       "    (tau): ReLU()\n",
       "    (V_d_transform): Identity()\n",
       "    (graph_transform): Identity()\n",
       "  )\n",
       "  (agg): MeanAggregation()\n",
       "  (bn): Identity()\n",
       "  (predictor): BinaryClassificationFFN(\n",
       "    (ffn): MLP(\n",
       "      (0): Sequential(\n",
       "        (0): Linear(in_features=300, out_features=300, bias=True)\n",
       "      )\n",
       "      (1): Sequential(\n",
       "        (0): ReLU()\n",
       "        (1): Dropout(p=0.0, inplace=False)\n",
       "        (2): Linear(in_features=300, out_features=4, bias=True)\n",
       "      )\n",
       "    )\n",
       "    (criterion): BCELoss(task_weights=[[1.0, 1.0, 1.0, 1.0]])\n",
       "    (output_transform): Identity()\n",
       "  )\n",
       "  (X_d_transform): Identity()\n",
       "  (metrics): ModuleList(\n",
       "    (0): BinaryAUROC()\n",
       "    (1): BCELoss(task_weights=[[1.0, 1.0, 1.0, 1.0]])\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mpnn = models.MPNN(mp, agg, ffn, batch_norm, metric_list)\n",
    "\n",
    "mpnn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set up trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: False, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "HPU available: False, using: 0 HPUs\n"
     ]
    }
   ],
   "source": [
    "trainer = pl.Trainer(\n",
    "    logger=False,\n",
    "    enable_checkpointing=True, # Use `True` if you want to save model checkpoints. The checkpoints will be saved in the `checkpoints` folder.\n",
    "    enable_progress_bar=True,\n",
    "    accelerator=\"cpu\",\n",
    "    devices=1,\n",
    "    max_epochs=20, # number of epochs to train for\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Start training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/knathan/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/callbacks/model_checkpoint.py:654: Checkpoint directory /home/knathan/chemprop/examples/checkpoints exists and is not empty.\n",
      "Loading `train_dataloader` to estimate number of stepping batches.\n",
      "/home/knathan/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:424: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.\n",
      "\n",
      "  | Name            | Type                    | Params | Mode \n",
      "--------------------------------------------------------------------\n",
      "0 | message_passing | BondMessagePassing      | 227 K  | train\n",
      "1 | agg             | MeanAggregation         | 0      | train\n",
      "2 | bn              | Identity                | 0      | train\n",
      "3 | predictor       | BinaryClassificationFFN | 91.5 K | train\n",
      "4 | X_d_transform   | Identity                | 0      | train\n",
      "5 | metrics         | ModuleList              | 0      | train\n",
      "--------------------------------------------------------------------\n",
      "319 K     Trainable params\n",
      "0         Non-trainable params\n",
      "319 K     Total params\n",
      "1.277     Total estimated model params size (MB)\n",
      "24        Modules in train mode\n",
      "0         Modules in eval mode\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sanity Checking DataLoader 0:   0%|          | 0/1 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/knathan/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:424: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 19: 100%|██████████| 7/7 [00:00<00:00,  9.07it/s, train_loss_step=0.238, val_loss=0.355, train_loss_epoch=0.345]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=20` reached.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 19: 100%|██████████| 7/7 [00:00<00:00,  8.90it/s, train_loss_step=0.238, val_loss=0.355, train_loss_epoch=0.345]\n"
     ]
    }
   ],
   "source": [
    "trainer.fit(mpnn, train_loader, val_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/knathan/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:424: The 'test_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing DataLoader 0: 100%|██████████| 1/1 [00:00<00:00, 40.45it/s]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\">        Test metric        </span>┃<span style=\"font-weight: bold\">       DataLoader 0        </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">         test/roc          </span>│<span style=\"color: #800080; text-decoration-color: #800080\">    0.6421189308166504     </span>│\n",
       "└───────────────────────────┴───────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1m       Test metric       \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m      DataLoader 0       \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│\u001b[36m \u001b[0m\u001b[36m        test/roc         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m   0.6421189308166504    \u001b[0m\u001b[35m \u001b[0m│\n",
       "└───────────────────────────┴───────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "results = trainer.test(mpnn, test_loader)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "chemprop",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
